import time

from robomaster import robot
from util import *


class drone_by_robomaster:
    def __init__(self):
        robot.config.LOCAL_IP_STR = "192.168.10.2"
        self.tl_drone = robot.Drone()
        self.tl_drone.initialize()
        self.battery = self.tl_drone.battery

        self.tl_camera = self.tl_drone.camera

        self.tl_flight = self.tl_drone.flight
        self.moving_speed = [20, 50, 75, 100]
        self.angular_speed = [20, 50, 75, 100]
        self.ascent_speed = [30, 60, 90, 100]
        self.gear = 0

    def drone_increase_gear(self):
        if self.gear < 3:
            self.gear += 1

    def drone_decrease_gear(self):
        if self.gear > 0:
            self.gear -= 1

    def drone_get_gear(self):
        return self.gear + 1

    def drone_get_battery(self):
        return self.battery.get_battery()

    @fn_timer
    def drone_move(self, a, b, c, d, speed_magnification=1):
        a_speed = -self.moving_speed[self.gear] * a * speed_magnification
        b_speed = self.moving_speed[self.gear] * b * speed_magnification
        c_speed = self.ascent_speed[self.gear] * c
        d_speed = -self.angular_speed[self.gear] * d
        a_speed = min(max(a_speed, -100), 100)
        b_speed = min(max(b_speed, -100), 100)
        c_speed = min(max(c_speed, -100), 100)
        d_speed = min(max(d_speed, -100), 100)
        log("drone_move", a, b, c, d, speed_magnification)
        self.tl_flight.rc(a=a_speed, b=b_speed, c=c_speed, d=d_speed)

    @fn_timer
    def drone_flip(self, direction):
        self.tl_flight.flip(direction)

    # 获取当前飞行速度
    def drone_get_speed(self):
        return self.tl_flight.get_speed()

    # 起飞
    @fn_timer
    def drone_takeoff(self):
        self.tl_flight.takeoff().wait_for_completed()

    # 停止
    @fn_timer
    def drone_stop(self):
        self.tl_flight.rc(a=0, b=0, c=0, d=0)

    # 降落
    @fn_timer
    def drone_land(self):
        self.tl_flight.land().wait_for_completed()

    def drone_init_camera(self):
        self.tl_camera.start_video_stream(display=False)
        self.tl_camera.set_fps("high")
        self.tl_camera.set_resolution("high")
        self.tl_camera.set_bitrate(6)

    def drone_get_img(self):
        try:
            img = self.tl_camera.read_cv2_image()
        except:
            img = None
        return img

    def drone_close_camera(self):
        self.tl_camera.stop_video_stream()

    def close(self):
        self.tl_camera.stop_video_stream()
        self.tl_drone.close()


if __name__ == '__main__':
    d = drone_by_robomaster()
    drone_version = d.tl_drone.get_sdk_version()
    log("Drone sdk version: {0}".format(drone_version))
    SN = d.tl_drone.get_sn()
    log("drone sn: {0}".format(SN))
